"""
数据集验证主流程

协调各个组件完成完整的验证流程
"""

import argparse
import logging
import os
import json
from typing import Dict, Any, Optional

# DataLoader 类已合并到此文件中
import random
from pre_labeler import PreLabeler
from validator import DatasetValidator
from result_analyzer import ResultAnalyzer
from config import Config

# 导入并设置统一的日志配置
from __init__ import setup_logging

# 设置日志系统
log_file = setup_logging("logs")
logger = logging.getLogger(__name__)
logger.info("数据集验证流程开始")


class DataLoader:
    """数据加载器"""
    
    def __init__(self):
        """初始化数据加载器"""
        logger.info("初始化数据加载器")
    
    def load_data(self, file_path: str) -> list:
        """
        加载验证数据
        
        Args:
            file_path: 数据文件路径
            
        Returns:
            数据列表
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"数据文件不存在: {file_path}")
        
        logger.info(f"加载数据文件: {file_path}")
        
        data = []
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    
                    try:
                        sample = json.loads(line)
                        data.append(sample)
                    except json.JSONDecodeError as e:
                        logger.warning(f"第{line_num}行JSON解析失败: {e}")
                        continue
            
            logger.info(f"成功加载 {len(data)} 条数据")
            return data
            
        except Exception as e:
            logger.error(f"加载数据文件失败: {e}")
            raise
    
    def get_random_samples(self, data: list, count: int = 5) -> list:
        """
        随机获取指定数量的样本
        
        Args:
            data: 数据列表
            count: 样本数量
            
        Returns:
            随机样本列表
        """
        random.seed(42)  # 设置随机种子以确保结果可重现
        
        if count >= len(data):
            return data
        
        return random.sample(data, count)


def load_config(config_path: str) -> Dict:
    if not config_path or not os.path.exists(config_path):
        return {}
    with open(config_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def run_pipeline(
    units_data_file: str,
    output_dir: str,
    sample_count: int,
    api_key: str = None,
    base_url: str = None,
    model_name: str = 'Qwen/Qwen3-8B',
    enable_pre_label: bool = True,
):

    # 1) 加载待验证数据
    loader = DataLoader()
    data = loader.load_data(units_data_file)
    if sample_count and sample_count < len(data):
        data = loader.get_random_samples(data, sample_count)
        logger.info(f"随机选择 {len(data)} 条样本用于验证")

    # 2) 可选：为样本做预标签，方便动态few-shot定位标签库
    if enable_pre_label and api_key and base_url:
        labeler = PreLabeler(api_key=api_key, base_url=base_url, model_name=model_name)
        logger.info("启用流式预标签：每条数据先预标签再验证")
    else:
        logger.info("跳过预标签阶段或缺少API配置")

    # 3) 执行验证并增量写入
    analyzer = ResultAnalyzer(output_dir)
    results_file = analyzer.get_new_results_file_path()

    validator = DatasetValidator(
        api_key=api_key or 'EMPTY',
        base_url=base_url or 'http://localhost:15070',
        model_name=model_name,
        label_data_dir=None,
    )

    for idx, sample in enumerate(data, 1):
        logger.info(f"验证样本 {idx}/{len(data)}: {sample.get('unit_id','unknown')}")
        
        # 流式预标签：先对当前样本做预标签
        if enable_pre_label and api_key and base_url:
            try:
                logger.info(f"为样本 {sample.get('unit_id','unknown')} 进行预标签...")
                is_valid, labeled_sample, skip_reason = labeler.process_single_data(sample)
                if is_valid:
                    # 预标签成功，使用带标签的样本进行验证
                    sample = labeled_sample
                    logger.info(f"样本 {sample.get('unit_id','unknown')} 预标签完成: {sample.get('pre_label', {}).get('label_name', 'N/A')}")
                else:
                    # 预标签失败或跳过，跳过后续验证任务
                    logger.info(f"样本 {sample.get('unit_id','unknown')} 跳过后续验证任务，原因: {skip_reason}")
                    continue  # 直接跳到下一条数据
            except Exception as e:
                logger.error(f"样本 {sample.get('unit_id','unknown')} 预标签失败: {e}")
                # 预标签失败时跳过后续验证任务
                logger.info(f"样本 {sample.get('unit_id','unknown')} 跳过后续验证任务，继续处理下一条数据")
                continue  # 直接跳到下一条数据
        
        try:
            sample_results = validator.validate_sample(sample)
            for _, result in sample_results.items():
                analyzer.append_validation_result(results_file, result)
        except Exception as e:
            logger.error(f"验证样本失败: {e}")
            continue

    # 5) 汇总与报告
    all_results = analyzer.load_results(results_file)
    analysis = analyzer.analyze_evaluation_results(all_results)
    analyzer.generate_analysis_report(analysis)
    analyzer.save_comparison_data(analyzer.create_comparison_data(all_results))

    logger.info("管道执行完成")


def parse_args():
    parser = argparse.ArgumentParser(description='文档验证主流程 pipeline')
    parser.add_argument('--config', type=str, default=None)
    parser.add_argument('--units-data-file', type=str, required=False,
                        default=Config.get_units_data_file())
    parser.add_argument('--output-dir', type=str, required=False,
                        default=Config.get_output_dir())
    parser.add_argument('--sample-count', type=int, default=Config.get_sample_count())
    parser.add_argument('--api-key', type=str, default=Config.get_api_key())
    parser.add_argument('--base-url', type=str, default=Config.get_base_url())
    parser.add_argument('--model-name', type=str, default=Config.get_model_name())
    parser.add_argument('--no-pre-label', action='store_true')
    return parser.parse_args()


def main():
    args = parse_args()
    cfg = load_config(args.config) if args.config else {}

    run_pipeline(
        units_data_file=cfg.get('data_file', args.units_data_file),
        output_dir=cfg.get('output_dir', args.output_dir),
        sample_count=int(cfg.get('sample_count', args.sample_count)),
        api_key=args.api_key or cfg.get('api_key'),
        base_url=args.base_url or cfg.get('base_url'),
        model_name=args.model_name or cfg.get('model_name', Config.get_model_name()),
        enable_pre_label=not args.no_pre_label,
    )


if __name__ == '__main__':
    main()


